#include <rclcpp/rclcpp.hpp>
#include <geometry_msgs/msg/wrench_stamped.hpp>
#include <moveit/move_group_interface/move_group_interface.h>
#include <chrono>
#include <cmath>
#include <fstream>
#include <iomanip>
#include <sstream>
#include <filesystem>
#include <thread>
#include <future>

using namespace std::chrono_literals;
namespace fs = std::filesystem;

class ForceMoveLogger : public rclcpp::Node
{
public:
  ForceMoveLogger()
  : Node("force_move_logger"),
    start_detected_(false),
    recording_(false),
    prev_F_(0.0)
  {
    // ---- 参数设置 ----
    delta_threshold_ = 0.1;          // 力变化起始阈值 (N)
    delay_after_motion_ms_ = 1500;   // 运动完成后延迟 (ms)

    // ---- 订阅力传感器 ----
    sub_wrench_ = this->create_subscription<geometry_msgs::msg::WrenchStamped>(
      "/wrench_comp", 10,
      std::bind(&ForceMoveLogger::wrenchCallback, this, std::placeholders::_1));

    RCLCPP_INFO(this->get_logger(), "节点初始化完成，等待开始运动...");
    timer_ = this->create_wall_timer(2s, std::bind(&ForceMoveLogger::startMotion, this));
  }

private:
  // ===== 参数 =====
  double delta_threshold_;
  int delay_after_motion_ms_;

  // ===== 状态变量 =====
  bool start_detected_;
  bool recording_;
  double prev_F_;

  struct ForceData {
    double time;
    double fx, fy, fz;
  };
  std::vector<ForceData> force_log_;

  rclcpp::Subscription<geometry_msgs::msg::WrenchStamped>::SharedPtr sub_wrench_;
  rclcpp::TimerBase::SharedPtr timer_;
  std::unique_ptr<moveit::planning_interface::MoveGroupInterface> move_group_;

  // ==================== 启动运动 ====================
  void startMotion()
  {
    timer_->cancel();

    move_group_ = std::make_unique<moveit::planning_interface::MoveGroupInterface>(
        shared_from_this(), "zc4");
    move_group_->setMaxVelocityScalingFactor(0.03);
    move_group_->setMaxAccelerationScalingFactor(0.03);

    // Step 1: 回到 middle 位姿
    RCLCPP_INFO(this->get_logger(), "移动到 middle 位姿...");
    move_group_->setNamedTarget("middle");
    move_group_->move();

    // Step 2: 设定目标位姿（你提供的）
    geometry_msgs::msg::Pose target_pose;
    target_pose.position.x = 0.0362074;
    target_pose.position.y = 0.0;
    target_pose.position.z = 0.0708187;
    target_pose.orientation.x = -0.000396547;
    target_pose.orientation.y = 0.220008;
    target_pose.orientation.z = 0.0;
    target_pose.orientation.w = 0.975498;

    RCLCPP_INFO(this->get_logger(), "开始移动到目标位姿...");

    // ✅ 异步执行运动，不阻塞回调
    std::thread([this, target_pose]() {
      move_group_->setPoseTarget(target_pose);
      move_group_->move();  // 阻塞在子线程中执行
      RCLCPP_INFO(this->get_logger(), "机械臂运动完成，延时 %d ms ...",
                  delay_after_motion_ms_);
      std::this_thread::sleep_for(std::chrono::milliseconds(delay_after_motion_ms_));
      saveCSV();
      rclcpp::shutdown();
    }).detach();
  }

  // ==================== 力回调 ====================
  void wrenchCallback(const geometry_msgs::msg::WrenchStamped::SharedPtr msg)
  {
    RCLCPP_INFO(this->get_logger(), "callback...");
    double Fx = msg->wrench.force.x;
    double Fy = msg->wrench.force.y;
    double Fz = msg->wrench.force.z;
    double F = std::sqrt(Fx*Fx + Fy*Fy + Fz*Fz);
    double t = msg->header.stamp.sec + msg->header.stamp.nanosec * 1e-9;

    // 检测力变化开始
    if (!start_detected_ && std::fabs(F - prev_F_) > delta_threshold_) {
      start_detected_ = true;
      recording_ = true;
      force_log_.clear();
      RCLCPP_INFO(this->get_logger(), "检测到力变化，开始记录...");
    }

    if (recording_) {
      force_log_.push_back({t, Fx, Fy, Fz});
    }

    prev_F_ = F;
  }

  // ==================== 保存CSV ====================
  void saveCSV()
  {
    if (force_log_.empty()) {
      RCLCPP_WARN(this->get_logger(), "未检测到力变化，无数据保存。");
      return;
    }

    fs::create_directories("force_logs");
    auto now = std::chrono::system_clock::now();
    std::time_t now_c = std::chrono::system_clock::to_time_t(now);
    std::stringstream ss;
    ss << "force_logs/force_segment_" 
       << std::put_time(std::localtime(&now_c), "%Y-%m-%d_%H%M%S") 
       << ".csv";
    std::string filename = ss.str();

    std::ofstream file(filename);
    file << "time,Fx,Fy,Fz\n";
    for (auto &d : force_log_)
      file << d.time << "," << d.fx << "," << d.fy << "," << d.fz << "\n";
    file.close();

    RCLCPP_INFO(this->get_logger(), "✅ 力数据已保存至 %s", filename.c_str());
  }
};

int main(int argc, char **argv)
{
  rclcpp::init(argc, argv);

  // ✅ 使用多线程执行器以保证订阅与控制并行
  rclcpp::executors::MultiThreadedExecutor executor;
  auto node = std::make_shared<ForceMoveLogger>();
  executor.add_node(node);
  executor.spin();

  rclcpp::shutdown();
  return 0;
}
